package mg2d.solver;

import mg2d.Matrix;
import mg2d.Tensor;
import mg2d.domain.Domain;
import mg2d.equation.ConvDiffEqn;

public class FaceCentered extends Multigrid {

	public FaceCentered(ConvDiffEqn cde, Domain domain, int n) {
		super(cde, domain, n);
	}
	
	private double plower(int i, int j, int level) {
		return (Math.sqrt(3.0)/2.0)*q(i,j,level)-(3.0/2.0)*p(i,j,level);
	}
	
	private double qlower(int i, int j, int level) {
		return -(Math.sqrt(3.0)/2.0)*q(i,j,level)-(3.0/2.0)*p(i,j,level);
	}
	
	private double flower(int i, int j, int level) {
		return (3.0/2.0)*f(i,j,level);
	}

	@Override
	void discret(int levels) {
		/*
		 * Calculate h value.
		 */
		this.h = 2.0*Math.sqrt(3.0)/(3.0*n);
		/*
		 * Initialize f,u,r,coeff
		 */
		f = new Matrix[levels];
		u = new Matrix[levels];
		r = new Matrix[levels];
		coeff = new Tensor[levels];
		for(int i=0; i<levels; i++) {
			f[i] = new Matrix(levelN(i)+1);
			u[i] = new Matrix(levelN(i)+1);
			r[i] = new Matrix(levelN(i)+1);
			coeff[i] = new Tensor(levelN(i)+1, 7);
		}
		
		/*
		 * Setup boundary of u.
		 */
		for(int i=0; i<=n; i++) {
			u[0].m[0][i] = u(0,i,0);
			u[0].m[n][i] = u(n,i,0);
			u[0].m[i][0] = u(i,0,0);
			u[0].m[i][n] = u(i,n,0);
		}
		
		/*
		 * Setup the values of coeff (all levels)
		 * Setup the values of f (fine level only)
		 */
		for(int level=0; level<levels; level++) {
			double hh = this.h * Math.pow(2.0, level);
			for(int i=1; i<=levelN(level)-1; i++) {
				for(int j=1; j<=levelN(level)-1; j++) {
					final double p = plower(i,j,level);
					final double q = qlower(i,j,level);
					final double f = flower(i,j,level);
					
					final double ps = (plower(i,j+1,level) - plower(i,j-1,level))/(2.0*hh);
					final double qs = (qlower(i,j+1,level) - qlower(i,j-1,level))/(2.0*hh);
					final double fs = (flower(i,j+1,level) - flower(i,j-1,level))/(2.0*hh);
					final double pr = (plower(i+1,j,level) - plower(i-1,j,level))/(2.0*hh);
					final double qr = (qlower(i+1,j,level) - qlower(i-1,j,level))/(2.0*hh);
					final double pt = (plower(i-1,j-1,level)- plower(i+1,j+1,level))/(2.0*hh);
					final double qt = (qlower(i-1,j-1,level)- qlower(i+1,j+1,level))/(2.0*hh);
					final double ft = (flower(i-1,j-1,level)- flower(i+1,j+1,level))/(2.0*hh);
					
					final double pss = (plower(i,j+1,level)-2.0*plower(i,j,level)+plower(i,j-1,level))/(hh*hh);
					final double qss = (qlower(i,j+1,level)-2.0*qlower(i,j,level)+qlower(i,j-1,level))/(hh*hh);
					final double prr = (plower(i+1,j,level)-2.0*plower(i,j,level)+plower(i-1,j,level))/(hh*hh);
					final double qrr = (qlower(i+1,j,level)-2.0*qlower(i,j,level)+qlower(i-1,j,level))/(hh*hh);
					final double ptt = (plower(i-1,j-1,level)-2.0*plower(i,j,level)+plower(i+1,j+1,level))/(hh*hh);
					final double qtt = (qlower(i-1,j-1,level)-2.0*qlower(i,j,level)+qlower(i+1,j+1,level))/(hh*hh);;
					
					final double A = hh*hh/24.0*( p*qt+q*pt - p*p*q - 2.0*q*pt - (prr + pss + ptt) );
					final double B = hh*hh/24.0*( p*qs+q*ps - q*q*p - 2.0*p*qs - (qrr + qss + qtt) );
					final double C = -hh*hh/12.0*( p*q + pt + qs );
					final double D = -hh*hh/12.0*( p*p/2.0 + ps + qr );
					final double E = -hh*hh/12.0*( q*q/2.0 + pr + qt );
					final double G = hh/6.0* ( p + q );
					final double H = (hh*hh*hh*hh)/24.0*( p*q*f + 2.0*p*fs + 2.0*q*ft - (p*fs+f*ps) - (q*ft+qt*f) );
					
					coeff[level].t[i][j][0] = -(6.0 - 2.0*C - 2.0*D - 2.0*E);
					coeff[level].t[i][j][1] = 1.0 - C - G;
					coeff[level].t[i][j][2] = 1.0 - q*hh/2.0 + B*hh/2.0 - E + G;
					coeff[level].t[i][j][3] = 1.0 + p*hh/2.0 - A*hh/2.0 - D - G;
					coeff[level].t[i][j][4] = 1.0 - C + G;
					coeff[level].t[i][j][5] = 1.0 + q*hh/2.0 - B*hh/2.0 - E - G; 
					coeff[level].t[i][j][6] = 1.0 - p*hh/2.0 + A*hh/2.0 - D + G;

					if(level == 0) {
						this.f[level].m[i][j] = (hh*hh/24.0)*(18.0*f
								+flower(i+1,j,level)+flower(i+1,j+1,level)+flower(i,j+1,level)
							    +flower(i-1,j,level)+flower(i-1,j-1,level)+flower(i,j-1,level)) + H;
					}
				}
			}
		}
	}

	@Override
	void restriction(int fineLevel, int coarseLevel) {
		int fineN = levelN(fineLevel);
		for(int i=1; i<=fineN-1; i++) {
			for(int j=1; j<=fineN-1; j++) {
				r[fineLevel].m[i][j] = f[fineLevel].m[i][j] - 
				   (  coeff[fineLevel].t[i][j][0]* u[fineLevel].m[i][j] 
                    + coeff[fineLevel].t[i][j][1]* u[fineLevel].m[i+1][j]   + coeff[fineLevel].t[i][j][2]* u[fineLevel].m[i+1][j+1] 
				    + coeff[fineLevel].t[i][j][3]* u[fineLevel].m[i][j+1]   + coeff[fineLevel].t[i][j][4]* u[fineLevel].m[i-1][j]
                    + coeff[fineLevel].t[i][j][5]* u[fineLevel].m[i-1][j-1] + coeff[fineLevel].t[i][j][6]* u[fineLevel].m[i][j-1]);
			}
		}
		
		int coarseN = levelN(coarseLevel);
		for(int i=1; i<=coarseN-1; i++) {
			final int fi = 2*i;
			for(int j=1; j<=coarseN-1; j++) {
				final int fj = 2*j;
				u[coarseLevel].m[i][j] = 0d;
				
				f[coarseLevel].m[i][j] = (4.0/8.0)*(2.0*r[fineLevel].m[fi][fj]
				    + r[fineLevel].m[fi+1][fj]   + r[fineLevel].m[fi+1][fj+1]
				    + r[fineLevel].m[fi][fj+1]   + r[fineLevel].m[fi-1][fj] 
				    + r[fineLevel].m[fi-1][fj-1] + r[fineLevel].m[fi][fj-1]);
			}
		}
	}

	@Override
	void prolongation(int coarseLevel, int fineLevel) {
		int coarseN = n/(int)Math.pow(2, coarseLevel);
		int fi, fj;
		for(int i=1; i<=coarseN-1; i++) {
			fi = i*2;
			for(int j=1; j<=coarseN-1; j++) {
				fj = j*2;
				
				u[fineLevel].m[fi][fj] += u[coarseLevel].m[i][j];
				u[fineLevel].m[fi+1][fj] += (.5)*u[coarseLevel].m[i][j];
				u[fineLevel].m[fi+1][fj+1] += (.5)*u[coarseLevel].m[i][j];
				u[fineLevel].m[fi][fj+1] += (.5)*u[coarseLevel].m[i][j];
				u[fineLevel].m[fi-1][fj] += (.5)*u[coarseLevel].m[i][j];
				u[fineLevel].m[fi-1][fj-1] += (.5)*u[coarseLevel].m[i][j];
				u[fineLevel].m[fi][fj-1] += (.5)*u[coarseLevel].m[i][j];
			}
		}
	}

	@Override
	double relax(int level) {
		int levelN = levelN(level);
		double unew;
		double residual = 0;
		
		for(int i=1; i<=levelN-1; i++) {
			for(int j=1; j<=levelN-1; j++) {
				unew = ( 
				  coeff[level].t[i][j][1]* u[level].m[i+1][j]   + coeff[level].t[i][j][2]* u[level].m[i+1][j+1] 
				+ coeff[level].t[i][j][3]* u[level].m[i][j+1]   + coeff[level].t[i][j][4]* u[level].m[i-1][j]
				+ coeff[level].t[i][j][5]* u[level].m[i-1][j-1] + coeff[level].t[i][j][6]* u[level].m[i][j-1] 
				- f[level].m[i][j] ) / ( -coeff[level].t[i][j][0] );
				
				residual += (u[level].m[i][j] - unew)*(u[level].m[i][j] - unew);
				u[level].m[i][j] = unew;
			}
		}
		return Math.sqrt(residual);	
	}
	
	public void extrapolate(Multigrid coarseGridSolution) {
		/*
		 * Extrapolate (even, even) points
		 */
		for(int i=1; i<coarseGridSolution.n-1; i++) {
			for(int j=1; j<coarseGridSolution.n-1; j++) {
				u[0].m[2*i][2*j] = (16.0*u[0].m[2*i][2*j] - coarseGridSolution.u[0].m[i][j]) / 15.0 ;
			}
		}
		/*
		 * Smooth the remaining points
		 */
		double residual = 0d;
		double ratio=1d;
		double r0 = 0d;
		int iter = 0;
		while(ratio > 10e-10) {
			residual = 0d;
			for(int i=1; i<=coarseGridSolution.n-1; i++) {
				for(int j=1; j<=coarseGridSolution.n-1; j++) {
					double unew;
					int fi = 2*i-1;
					int fj = 2*j-1;
					unew = ( 
						  coeff[0].t[fi][fj][1]* u[0].m[fi+1][fj]   + coeff[0].t[fi][fj][2]* u[0].m[fi+1][fj+1] 
						+ coeff[0].t[fi][fj][3]* u[0].m[fi][fj+1]   + coeff[0].t[fi][fj][4]* u[0].m[fi-1][fj]
						+ coeff[0].t[fi][fj][5]* u[0].m[fi-1][fj-1] + coeff[0].t[fi][fj][6]* u[0].m[fi][fj-1] 
						- f[0].m[fi][fj] ) / ( -coeff[0].t[fi][fj][0] );
							
					residual += (u[0].m[fi][fj] - unew)*(u[0].m[fi][fj] - unew);
					u[0].m[fi][fj] = unew;
				}
			}
			for(int i=1; i<=coarseGridSolution.n-1; i++) {
				for(int j=1; j<=coarseGridSolution.n-1; j++) {
					int fi = 2*i-1;
					int fj = 2*j;
					double unew = ( 
						  coeff[0].t[fi][fj][1]* u[0].m[fi+1][fj]   + coeff[0].t[fi][fj][2]* u[0].m[fi+1][fj+1] 
						+ coeff[0].t[fi][fj][3]* u[0].m[fi][fj+1]   + coeff[0].t[fi][fj][4]* u[0].m[fi-1][fj]
						+ coeff[0].t[fi][fj][5]* u[0].m[fi-1][fj-1] + coeff[0].t[fi][fj][6]* u[0].m[fi][fj-1] 
						- f[0].m[fi][fj] ) / ( -coeff[0].t[fi][fj][0] );
								
					residual += (u[0].m[fi][fj] - unew)*(u[0].m[fi][fj] - unew);
					u[0].m[fi][fj] = unew;
				}
			}
			for(int i=1; i<=coarseGridSolution.n-1; i++) {
				for(int j=1; j<=coarseGridSolution.n-1; j++) {
					int fi = 2*i;
					int fj = 2*j-1;
					double unew = ( 
						  coeff[0].t[fi][fj][1]* u[0].m[fi+1][fj]   + coeff[0].t[fi][fj][2]* u[0].m[fi+1][fj+1] 
						+ coeff[0].t[fi][fj][3]* u[0].m[fi][fj+1]   + coeff[0].t[fi][fj][4]* u[0].m[fi-1][fj]
						+ coeff[0].t[fi][fj][5]* u[0].m[fi-1][fj-1] + coeff[0].t[fi][fj][6]* u[0].m[fi][fj-1] 
						- f[0].m[fi][fj] ) / ( -coeff[0].t[fi][fj][0] );
								
					residual += (u[0].m[fi][fj] - unew)*(u[0].m[fi][fj] - unew);
					u[0].m[fi][fj] = unew;
				}
			}
			if(iter == 0) {
				r0 = Math.sqrt(residual);
			} else {
				ratio = Math.sqrt(residual)/r0;
			}
			iter++;
		}
		System.out.println("Extrapolation Iterations: "+iter);
	}
	
	public double getEvenEvenError() {
		double maxError = 0;
		double error;
		for(int i=0; i<=n; i+=2) {
			for(int j=0; j<=n; j+=2) {
				error = Math.abs(u[0].m[i][j] - u(i,j,0));
				if(error > maxError) {
					maxError = error;
				}
			}
		}
		return maxError;
	}

	@Override
	double coarseRelax() {
		return relax( log2(n)-1 );
	}

	@Override
	double residual2norm() {
		double norm = 0;
		for(int i=1; i<=n-1; i++) {
			for(int j=1; j<=n-1; j++) {
				norm += r[0].m[i][j]*r[0].m[i][j];
			}
		}
		return norm;
	}
}
